import torch.nn as nn
from utils import *
from scipy import sparse
from scipy.sparse import coo_matrix
from torch.nn import Module


class GNNLayer(Module):
    def __init__(self, args, inF, outF):
        super(GNNLayer, self).__init__()
        self.args = args
        self.device = self.args.device
        self.inF = inF
        self.outF = outF
        self.linear = torch.nn.Linear(in_features=inF, out_features=outF)
        self.interActTransform = torch.nn.Linear(
            in_features=inF, out_features=outF)

    def forward(self, laplacianMat, selfLoop, features):
        # for GCF ajdMat is a (N+M) by (N+M) mat
        # laplacianMat L = D^-1(A)D^-1 # 拉普拉斯矩阵
        L1 = laplacianMat + selfLoop
        L2 = laplacianMat
        L1 = L1
        inter_feature = torch.sparse.mm(L2, features)  # 矩阵乘法
        inter_feature = torch.mul(inter_feature, features)  # 对应位置相乘

        inter_part1 = self.linear(torch.sparse.mm(L1, features))
        inter_part2 = self.interActTransform(
            torch.sparse.mm(L2, inter_feature))

        return inter_part1 + inter_part2


class GCF(Module):
    def __init__(self, args, user_feature, item_feature, rating):
        super(GCF, self).__init__()
        self.args = args
        self.device = args.device
        self.user_feature = user_feature
        self.item_feature = item_feature
        self.rating = rating
        self.userNum = rating['userId'].max() + 1
        self.itemNum = rating['itemId'].max() + 1
        # age:60 gender:2 occupation:20 zip-code:794
        self.user_age_embedding = nn.Embedding(user_feature['age'].max() + 1, 10)
        self.user_gender_embedding = nn.Embedding(user_feature['gender'].max() + 1, 2)
        self.user_occupation_embedding = nn.Embedding(user_feature['occupation'].max() + 1, 4)
        self.user_zipCode_embedding = nn.Embedding(user_feature['location'].max() + 1, 48)
        # summary:24 Wind Bearing (degrees):7
        self.item_summary_embedding = nn.Embedding(item_feature['Summary'].max() + 1, 5)
        self.item_wind_embedding = nn.Embedding(item_feature['Wind Bearing (degrees)'].max() + 1, 3)
        self.item_feature_transform = nn.Linear(5, 56)

        self.GNNlayers = torch.nn.ModuleList()
        self.LaplacianMat = self.buildLaplacianMat(rating)  # sparse format
        self.leakyRelu = nn.LeakyReLU()
        self.selfLoop = self.getSparseEye(self.userNum + self.itemNum)
        self.transForm = nn.Linear(in_features=self.args.embedSize * (self.args.gcn_layers + 1),
                                   out_features=self.args.embedSize)

        for _ in range(self.args.gcn_layers):
            self.GNNlayers.append(GNNLayer(self.args, self.args.embedSize, self.args.embedSize))

    def getSparseEye(self, num):
        i = torch.LongTensor(
            [[k for k in range(0, num)], [j for j in range(0, num)]])
        val = torch.FloatTensor([1] * num)
        return torch.sparse.FloatTensor(i, val).to(self.device)

    def buildLaplacianMat(self, data):
        """构造拉普拉斯矩阵"""
        rt_item = data['itemId'] + self.userNum
        uiMat = coo_matrix((data['rating'], (data['userId'], data['itemId'])))

        uiMat_upperPart = coo_matrix((data['rating'], (data['userId'], rt_item)))
        uiMat = uiMat.transpose()
        uiMat.resize((self.itemNum, self.userNum + self.itemNum))

        A = sparse.vstack([uiMat_upperPart, uiMat])
        sumArr = (A > 0).sum(axis=1)  # [[sum],[sum]]
        diag = list(np.array(sumArr.flatten())[0])
        diag = np.power(diag, -0.5)
        D = sparse.diags(diag)  # 归一化的参数
        L = D * A * D  # 这里是矩阵乘法，不是对应位置相乘
        L = sparse.coo_matrix(L)
        row = L.row
        col = L.col
        i = torch.LongTensor([row, col])
        data = torch.FloatTensor(L.data)
        SparseL = torch.sparse.FloatTensor(i, data)
        return SparseL.to(self.device)

    def getFeatureMat(self):
        age = self.user_age_embedding(torch.tensor(self.user_feature['age']).to(self.device))
        gender = self.user_gender_embedding(torch.tensor(self.user_feature['gender']).to(self.device))
        occupation = self.user_occupation_embedding(torch.tensor(self.user_feature['occupation']).to(self.device))
        zipCode = self.user_zipCode_embedding(torch.tensor(self.user_feature['location']).to(self.device))
        userEmbd = torch.cat((age,
                              gender,
                              occupation,
                              zipCode), dim=1)
        summary = self.item_summary_embedding(torch.tensor(self.item_feature['Summary']).to(self.device))
        wind = self.item_wind_embedding(torch.tensor(self.item_feature['Wind Bearing (degrees)']).to(self.device))
        feature_tmp = torch.FloatTensor(self.item_feature.drop(
            columns=['itemId', 'Summary', 'Wind Bearing (degrees)']).to_numpy()).to(self.device)
        feature_transform = self.item_feature_transform(feature_tmp)
        itemEmbd = torch.cat((summary,
                              wind,
                              feature_transform), dim=1)
        features = torch.cat([userEmbd, itemEmbd], dim=0)
        return features.to(self.device)

    def forward(self, userIdx, itemIdx):
        itemIdx = itemIdx + self.userNum
        userIdx = list(userIdx.cpu().data)
        itemIdx = list(itemIdx.cpu().data)
        # gcf data propagation
        features = self.getFeatureMat()
        finalEmbed = features.clone()
        # 拼接
        for gcn in self.GNNlayers:
            features = gcn(self.LaplacianMat, self.selfLoop, features)
            features = nn.ReLU()(features)
            finalEmbed = torch.cat((finalEmbed, features.clone()), dim=1)
        userEmbed = finalEmbed[userIdx]
        itemEmbed = finalEmbed[itemIdx]
        userEmbed = self.transForm(userEmbed)
        itemEmbed = self.transForm(itemEmbed)
        embed = torch.mul(userEmbed, itemEmbed).sum(1)
        prediction = embed.flatten()
        return prediction
